"""
文档加载，并按一定条件切割成片段
将切割的文本片段灌入检索引擎
封装检索接口
构建调用流程：Query -> 检索 -> Prompt -> LLM -> 回复
"""
from ragbot import RAG_Bot
from extract_pdf import extract_text_from_pdf
from openai import OpenAI
import os
from dotenv import load_dotenv,find_dotenv
from vectorDB import MyVectorDBConnector

_ =load_dotenv(find_dotenv())

client = OpenAI(
    api_key = os.getenv("OPENAI_API_KEY"),
    base_url = os.getenv("OPENAI_API_BASE")


)


def get_completion(prompt, model="gpt-3.5-turbo"):
    '''封装 openai 接口'''
    messages = [{"role": "user", "content": prompt}]
    response = client.chat.completions.create(
        model=model,
        messages=messages,
        temperature=0,  # 模型输出的随机性，0 表示随机性最小
    )
    return response.choices[0].message.content


paragraphs = extract_text_from_pdf(r"/Users/SITA/JiWenBo/work-cood/zhihu/RAG/rzf.pdf",page_numbers=[0,1],min_line_length=10)
def get_embeddings(texts, model="text-embedding-ada-002"):
    '''封装 OpenAI 的 Embedding 模型接口'''
    data = client.embeddings.create(input = texts, model=model).data
    return [x.embedding for x in data]


"""

主流程


"""

vector_db = MyVectorDBConnector("pdf",get_embeddings)
# 向向量数据库中添加文档
vector_db.add_document(paragraphs)

user_query = "思科为什么起诉华为"

results = vector_db.search(user_query, 2)

for para in results['documents'][0]:
    print(para+'\n')    



"""创建一个 bot 机器人"""

bot = RAG_Bot(
    vector_db,
    llm_api=get_completion
)
user_query = "任正非一共在几家公司待过"
response = bot.chat(user_query)
print(response)